#include <opencv2/opencv.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <iostream>
#include <complex>

#include "split_bregman.h"
#include "cvabs.h"
#include "elementwise.h"

using namespace cv;
using namespace std;

typedef complex<float> usetype;

// ============================================================
// = Function declarations                                    =
// ============================================================
void get_input_mat( string filein, Mat* imout );
void get_fourier_mask( int rows, int cols, Mat* R, float sampleprob );

// ============================================================
// = main                                                     =
// ============================================================

int main()
{
	Mat_<usetype> imin, imF, imRF, imFabs, r, R, u;
	Mat uabs;
	const float mu = 10.0;
	const float lambda = mu;
	const float gamma =  mu/100.0;
	const int nInner = 20;
	const int nOuter = 10;

	// setup display windows
	namedWindow("Input image", 1);
	namedWindow("Fourier magnitude",1);
	namedWindow("Reconstruction", 1);

	// Create input data imin
	get_input_mat("pr3ut.png", &imin);
	// create Fourier data imF
	dft(imin, imF);
	imF = imF/double(imin.rows);
	// create abs of imF for later viewing purposes, scale it
	cvabs(&imF, &imFabs); 
	imFabs = imFabs;
	// create Fourier mask
	get_fourier_mask(imin.rows, imin.cols, &R, 0.9);
	// create measurement data (i.e. undersampled Fourier data)
	imRF = elmult(&R, &imF);

	// display initial images
	//imshow("Input image", imin);
	vector<Mat> iminplanes;
	split(imin, iminplanes);
	imshow("Input image", iminplanes[0]);
	imshow("Fourier magnitude", imFabs);

	// Make reconstruction u via split_bregman
	split_bregman(&R, &imRF, &u, mu, lambda, gamma, nInner, nOuter);

	// take absolute value, find min and max for debugging
	cvabs(&u, &uabs);
	//double minel, maxel;
	//minMaxIdx(uabs, &minel, &maxel);

	// display reconstruction
	imshow("Reconstruction", uabs);
	waitKey(0);
	return 0;
}

// ============================================================
// = Auxiliary functions                                      =
// ============================================================

void get_input_mat( string filein, Mat* imout )
// get_input_mat: Reads in image at filein, converts to a normalized CV_32FC2 Mat at *imout
{
	Mat imin, iminbw;
	imin = imread(filein);
	cvtColor(imin, iminbw, CV_RGB2GRAY, 1);
	iminbw.convertTo(iminbw, CV_32F);
	iminbw *= 1./255;

	Mat planes[] = {iminbw, Mat::zeros(imin.size(), CV_32F)};
	merge(planes, 2, *imout);
}

void get_fourier_mask( int rows, int cols, Mat* R, float sampleprob = 0.5)
// get_fourier_mask: Creates a CV_32FC2 matrix at *R of size {rows, cols}. Each entry of *R is 1.0 with probability sampleprob, 0 otherwise.
{
	Mat initr = Mat(rows, cols, CV_32FC1, 0.0); // initial random data for mask generation
	Mat r; // single-channel mask
	*R = Mat_<usetype>(rows, cols, 0.0); // dual-channel mask, 2nd channel (i.e. imaginary part) all zeros

	RNG rng;
	rng.fill(initr, RNG::UNIFORM, 0, 1);
	compare(initr, sampleprob, r, CMP_LT);
	r.convertTo(r, CV_32F);
	r *= 1./255;
	r.at<float>(0) = 1; // mask needs to sample at the origin in Fourier space.
	Mat planes[] = {r, Mat::zeros(rows, cols, CV_32F)};
	merge(planes, 2, *R);
}